#include "intel_itt_wrapper.hpp"
#include "matrix_util.hpp"
#include "mpi/cipher_cannon.hpp"
#include <gtest/gtest.h>

#define COMM_WORLD_ROOT 0

INTEL_ITT_DOMAIN_CREATE(hhmm2d_domain, "test.mpi.hhmm2d");
INTEL_ITT_STRING_HANDLE_CREATE(hhmm2d_handle, "HMM2D");

params<helib::BGV> my_params;
bool check = false;
unsigned long input_bit_len;

#define PCGEMM_CANNON_TEST(NUM_P, DIM_M, DIM_N, DIM_K)                         \
  TEST(TEST_P##NUM_P, M##DIM_M##_N##DIM_N##_K##DIM_K) {                        \
    int rank, np;                                                              \
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);                                      \
    MPI_Comm_size(MPI_COMM_WORLD, &np);                                        \
    int sqrt_np = static_cast<int>(std::sqrt(np));                             \
    np = sqrt_np * sqrt_np;                                                    \
    if (np != NUM_P) {                                                         \
      GTEST_SKIP() << "The parameter of mpirun's `np` option do not meet the " \
                      "requirements of the current test suite";                \
    }                                                                          \
                                                                               \
    unsigned long m = DIM_M;                                                   \
    unsigned long n = DIM_N;                                                   \
    unsigned long k = DIM_K;                                                   \
    hmm_status<shmm_engine> status(METHOD::BSGS);                              \
    shhmm_cannon cce(COMM_WORLD_ROOT, MPI_COMM_WORLD, m, n, k, my_params,      \
                     status);                                                  \
                                                                               \
    int *A = nullptr;                                                          \
    int *B = nullptr;                                                          \
    int *C = nullptr;                                                          \
    if (rank == COMM_WORLD_ROOT) {                                             \
      FIX_RAND_SEED(true);                                                     \
      const std::size_t size_A = m * k;                                        \
      const std::size_t size_B = k * n;                                        \
      ALLOC_INT_ARRAY_WITHOUT_INIT(A, size_A);                                 \
      ALLOC_INT_ARRAY_WITHOUT_INIT(B, size_B);                                 \
      RAND_ARRAY(A, size_A, input_bit_len);                                    \
      RAND_ARRAY(B, size_B, input_bit_len);                                    \
      if (check) {                                                             \
        const std::size_t size_C = m * n;                                      \
        ALLOC_INT_ARRAY_WITHOUT_INIT(C, size_C);                               \
        cblas_igemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, 1.0,   \
                    A, k, B, n, 0.0, C, n);                                    \
      }                                                                        \
    }                                                                          \
                                                                               \
    /* Distribute Enc(input) */                                                \
    cce.init_data(A, B);                                                       \
                                                                               \
    /* Execute Cannon's algorithm & Get partial Enc(result) */                 \
    cce.hmme_status()->threads(omp_get_max_threads());                         \
    std::string partial_C;                                                     \
    INTEL_ITT_RESUME;                                                          \
    INTEL_ITT_TASK_BEGIN(hhmm2d_domain, hhmm2d_handle);                        \
    INTEL_ITT_TASK_END(hhmm2d_domain);                                         \
    INTEL_ITT_DETACH;                                                          \
    cce.exec(partial_C);                                                       \
    if (rank == COMM_WORLD_ROOT) {                                             \
      std::stringstream ss;                                                    \
      ss << "* CGEMM Method          : "                                       \
         << shmm_engine::METHOD_STR[cce.hmme_status()->method()] << std::endl; \
      ss << "* CGEMM # of threads    : " << cce.hmme_status()->threads()       \
         << std::endl;                                                         \
      cce.report(ss);                                                          \
      std::cout << ss.str();                                                   \
      std::flush(std::cout);                                                   \
      /*std::ofstream ss("cannon-report.txt");                                 \
      if (ss.is_open()) {                                                      \
        ss << "* CGEMM Method          : "                                     \
           << shmm_engine::METHOD_STR[cce.hmme_status()->method()]             \
           << std::endl;                                                       \
        ss << "* CGEMM # of threads    : " << cce.hmme_status()->threads()     \
           << std::endl;                                                       \
        cce.report(ss);                                                        \
        ss.close();                                                            \
      } else {                                                                 \
        std::cerr << "Unable to open the file 'report.txt' for writing."       \
                  << std::endl;                                                \
        MPI_Abort(MPI_COMM_WORLD, 1);                                          \
      }*/                                                                      \
    }                                                                          \
                                                                               \
    /* Only the root will get a vector of Dec(Enc(result)) */                  \
    std::vector<NTL::mat_ZZ> res;                                              \
    cce.collect_data(partial_C, res);                                          \
                                                                               \
    if (rank != COMM_WORLD_ROOT) {                                             \
      return;                                                                  \
    }                                                                          \
    if (!check) {                                                              \
      delete[] A;                                                              \
      delete[] B;                                                              \
      GTEST_SKIP() << "Skip checking answers";                                 \
    }                                                                          \
    /* Check answers */                                                        \
    const auto mb = cce.block_m();                                             \
    const auto nb = cce.block_n();                                             \
    for (std::size_t k = 0; k < np; k++) {                                     \
      const std::size_t i = k / sqrt_np;                                       \
      const std::size_t j = k % sqrt_np;                                       \
      for (std::size_t ii = 0; ii < mb; ii++) {                                \
        for (std::size_t jj = 0; jj < nb; jj++) {                              \
          auto expect = C[(i * mb + ii) * n + j * nb + jj];                    \
          auto actual = res[k][ii][jj];                                        \
          ASSERT_EQ(expect, actual) << "diff at (" << i << ", " << j << ")";   \
        }                                                                      \
      }                                                                        \
    }                                                                          \
  }

PCGEMM_CANNON_TEST(4, 4, 4, 4)
PCGEMM_CANNON_TEST(4, 6, 6, 6)
PCGEMM_CANNON_TEST(9, 6, 6, 6)

PCGEMM_CANNON_TEST(16, 256, 256, 256)
PCGEMM_CANNON_TEST(16, 512, 512, 512)
PCGEMM_CANNON_TEST(16, 1024, 1024, 1024)

PCGEMM_CANNON_TEST(64, 512, 512, 512)
PCGEMM_CANNON_TEST(64, 1024, 1024, 1024)

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);
  int ret = MPI_Init(&argc, &argv);
  EXPECT_EQ(ret, 0);

  // Get MPI metadata
  int rank;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);

  // Only one printer
  ::testing::TestEventListeners &listeners =
      ::testing::UnitTest::GetInstance()->listeners();
  if (rank != COMM_WORLD_ROOT) {
    //* Only the root has a secret key
    delete listeners.Release(listeners.default_result_printer());
  }

  // Parse command line arguments
  command_line_parser parser(my_params);
  parser.add_option("-chk", "enable_check_answer", &check);
  parser.add_option("-ib", "7", "input_bit_length", &input_bit_len);
  ret = parser.parse(argc, argv, rank == COMM_WORLD_ROOT);
  EXPECT_EQ(ret, 0);

  ret = RUN_ALL_TESTS();
  EXPECT_EQ(ret, 0);

  MPI_Finalize();
  return 0;
}
